import os
import random

# 设置路径
data_root = "./data/ECSSD/test"
image_dir = os.path.join(data_root, "Imgs")
mask_dir = os.path.join(data_root, "mask")
lst_path = os.path.join(data_root, "train_pair.lst")

# 获取所有图像文件
image_files = [f for f in os.listdir(image_dir) if f.endswith(('.jpg', '.png'))]

# 创建图像-掩码对
pairs = []
for img in image_files:
    mask = os.path.splitext(img)[0] + ".png"
    pairs.append(f"Imgs/{img} mask/{mask}")

# 随机打乱（如果需要）
random.shuffle(pairs)

# 写入列表文件
with open(lst_path, "w") as f:
    f.write("\n".join(pairs))

print(f"成功创建 {len(pairs)} 个训练对: {lst_path}")